# -*- coding: utf-8 -*-

import numpy as np
import cv2
import tensorflow as tf
from segnet import seg_net_network


IMAGE_HEIGHT = 256
IMAGE_WIDTH = 480
IMAGE_DEPTH = 3
GROUND_TRUTH_DEPTH = 1
GPU_NUMBER = 1


def _read_image(image_path):
    image = cv2.imread(image_path)
    image = cv2.resize(image, dsize=(IMAGE_WIDTH, IMAGE_HEIGHT))
    image = image.astype(dtype=np.float32) - [116.779, 123.68, 103.939]
    assert isinstance(image, np.ndarray)
    image = np.expand_dims(image, axis=0)
    return image


def show_road_image(image_path, road_mask, name="color_map"):
    image = cv2.imread(image_path)

    min_value, max_value = np.min(road_mask), np.max(road_mask)
    img = ((road_mask - min_value) / max_value) * 255
    img = img.astype(np.uint8, copy=True)
    _, img = cv2.threshold(img, 127, 255, cv2.THRESH_BINARY)
    img = cv2.medianBlur(img, ksize=5)
    temp = np.zeros_like(image)
    temp[:, :, 0] = img
    temp[:, :, 1] = img
    temp[:, :, 2] = img
    # color_map = cv2.applyColorMap(img, cv2.COLORMAP_JET)

    show_image = cv2.addWeighted(image, 0.3, temp, 0.7, 0)
    cv2.namedWindow(name, cv2.WINDOW_NORMAL)
    cv2.imshow(name, show_image)
    cv2.waitKey()
    cv2.destroyWindow(name)
    cv2.imwrite("result.jpg", show_image)


def main():
    image_path = "/home/lijun/DataBase/road_segment/480_256/images/0000368.jpg"
    image = _read_image(image_path)
    with tf.Graph().as_default():
        input_x = tf.placeholder(dtype=tf.float32,
                                 shape=[1, IMAGE_HEIGHT, IMAGE_WIDTH, IMAGE_DEPTH])
        predict = seg_net_network(in_put=input_x, is_training=False)
        _, road_tensor = tf.split(value=tf.nn.softmax(predict), num_or_size_splits=2, axis=3)
        saver = tf.train.Saver()
        with tf.Session() as session:
            saver.restore(session, "./model/480_256/seg_net-60000")
            road_image = session.run(road_tensor, feed_dict={input_x: image})
            road_image = road_image[0][:, :, 0]
            show_road_image(image_path, road_image)


if __name__ == "__main__":
    main()
